Skip to content

feat: lower LogicalNot, AllClose and Phase blocks in the native Expr compiler - #192

Merged
ausimian merged 1 commit into
feat/expr-compilerfrom
feat/expr-compiler-misc-blocks
Jun 6, 2026
Merged

feat: lower LogicalNot, AllClose and Phase blocks in the native Expr compiler#192
ausimian merged 1 commit into
feat/expr-compilerfrom
feat/expr-compiler-misc-blocks

Conversation

@ausimian

@ausimian ausimian commented Jun 6, 2026

Copy link
Copy Markdown
Owner

Summary

Closes the misc Nx.Block cluster on #188. Three new lower_block clauses, no new opcodes — each routes through primitives the IR already lowers, bit-identical to the Evaluator.

  • Nx.Block.LogicalNot — emits the existing :logical_not opcode directly with a trailing coerce, mirroring Emily.Backend.native_logical_not/2.
  • Nx.Block.AllClose — composes five existing primitives end-to-end, matching Emily.Backend.native_all_close/4 op-for-op: cast both to the merged float type, compute abs(a - b) <= atol + rtol * abs(b), optionally OR with isnan(a) AND isnan(b) for equal_nan: true, then reduce over every axis via :all. The trailing coerce produces {:u, 8}.
  • Nx.Block.Phase — falls through to the block's composed expansion (atan2(imag(t), real(t))) via the TopK-style block-parameter seeding (block-local :parameter nodes are fresh and would otherwise resolve to the outer function's input slots). All three primitives in the expansion are lowered as of feat: lower 19 more unary ops in the native Expr compiler #189 / feat: lower atan2, quotient and logical_xor in the native Expr compiler #190, so the result is bit-identical to the Evaluator.

Probe (scripts/expr_op_coverage.exs) drops 12 → 9 misses; the misc Nx.Block section of #188 is now empty.

Test plan

  • mix precommit clean (718 tests, 79 properties, 40 doctests, 0 failures — 6 new equivalence tests).
  • Op-coverage probe shows all three moved from MISS → OK.
  • New equivalence tests in compiler_equivalence_test.exs:
    • logical_not on both integer (s32) and float (f32) inputs, asserting the {:u, 8} predicate dtype.
    • all_close on close inputs (default tolerance → 1), far inputs (→ 0), custom rtol/atol inputs covering boundary cases (atol = 0.2 → close, atol = 0.05 → far), and the equal_nan: true branch (NaN compares equal to NaN, false is the default and returns 0).
    • phase across the four complex quadrants plus the real and imaginary axes (phase == 0 for positive real, ±pi for negative real, ±pi/2 for imaginary).
  • No opcode changes — opcode parity stays at the current kOpcodeCount (115).

Follows the same pattern as #189 / #190 / #191.

…compiler

Closes the misc-block cluster on the Expr op-coverage checklist (#188).
Three new lower_block clauses, no new opcodes — each routes through
primitives the IR already lowers, bit-identical to the Evaluator.

- `Nx.Block.LogicalNot` emits the existing `:logical_not` opcode
  directly with a trailing coerce, mirroring
  Emily.Backend.native_logical_not/2.
- `Nx.Block.AllClose` composes five existing primitives end-to-end,
  matching Emily.Backend.native_all_close/4 op-for-op: cast both to
  the merged float type, compute `abs(a - b) <= atol + rtol * abs(b)`,
  optionally OR with `isnan(a) AND isnan(b)` for `equal_nan: true`,
  then reduce over every axis via `:all`. The trailing coerce produces
  the {:u, 8} predicate dtype.
- `Nx.Block.Phase` falls through to the block's composed expansion —
  `atan2(imag(t), real(t))` — using the TopK-style block-parameter
  seeding (block-local `:parameter` nodes are fresh and would otherwise
  resolve to the outer function's input slots). All three primitives in
  the expansion (atan2/imag/real) are now lowered, so the result is
  bit-identical to the Evaluator.

Probe drops from 12 → 9 misses. Tests cover logical_not (assertion on
{:u, 8} dtype), all_close on close/far inputs, custom rtol/atol, and
the equal_nan branch, and phase across the four complex quadrants plus
the real/imaginary axes.
@ausimian
ausimian merged commit 8c3b15f into feat/expr-compiler Jun 6, 2026
3 checks passed
@ausimian
ausimian deleted the feat/expr-compiler-misc-blocks branch June 6, 2026 11:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant